/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "mstep.h"
#include "qfunc.h"

#include <assert.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

#include "fastonebigheader.h"

#define exp fastexp
#define log fastlog

static void
make_index (unsigned int* index,
            unsigned int  n_ratings)
{
  index[0] = 0;

  for (unsigned int i = 1; i < n_ratings; ++i)
    {
      unsigned int j = (unsigned int) ((i + 1) * drand48 ());
      index[i] = index[j];
      index[j] = i;
    }
}

float
mstep_sgd (float*                     loga,
           unsigned int               n_workers,
           unsigned int               rank,
           float*                     gamma,
           float*                     logbeta,
           unsigned int               n_images,
           float*                     c,
           unsigned int               n_labels,
           float*                     d,
           const float*               pz,
           Rating*                    ratings,
           unsigned int               n_ratings,
           ParameterizedDistribution* priorloga,
           Distribution*              priorgamma,
           NormalDistribution*        priorlogbeta,
           Distribution*              priorc,
           Distribution*              priord,
           float                      eta)   /* e.g, eta = 1 */
{
  float dloga[n_workers * rank];
  float dgamma[rank];
  float dlogbeta[n_images];
  float dc[n_labels * rank];
  float dd[n_labels * rank];
  unsigned int index[n_ratings];
  unsigned int lastimage[n_images];

  memset (lastimage, 0, sizeof (lastimage));

  make_index (index, n_ratings);

  for (unsigned int n = 0; n < n_ratings; ++n)
    {
      float thiseta = eta / sqrt (1.0 + n);
      unsigned int image = ratings[index[n]].image;

      memset (dloga, 0, n_workers * rank * sizeof (dloga[0]));
      memset (dgamma, 0, rank * sizeof (dgamma[0]));
      dlogbeta[image] = 0;
      memset (dc, 0, rank * n_labels * sizeof (dc[0]));
      memset (dd, 0, rank * n_labels * sizeof (dd[0]));

      float decay = 
        exp (2.0 * 
             eta * 
             (sqrt (1.0 + lastimage[image]) - sqrt (1.0 + n)) /
             (n_ratings * priorlogbeta->sq_stddev));

      logbeta[image] = 
        decay * logbeta[image] + (1.0 - decay) * priorlogbeta->mean;

      dqfunc_single (index[n], dloga, dgamma, dlogbeta, dc, dd, loga, n_workers,
                     rank, gamma, logbeta, n_images, c, n_labels, d, pz, 
                     ratings, n_ratings, priorloga, priorgamma, 
                     &priorlogbeta->base, priorc, priord);

      for (unsigned int i = 0; i < n_workers * rank; ++i)
        {
          loga[i] += thiseta * dloga[i];
        }

      for (unsigned int i = 0; i < rank; ++i)
        {
          gamma[i] += thiseta * dgamma[i];
        }

      logbeta[image] += thiseta * dlogbeta[image];
      lastimage[image] = 1 + n;

      for (unsigned int i = 0; i < n_labels * rank; ++i)
        {
          c[i] += thiseta * dc[i];
          d[i] += thiseta * dd[i];
        }
    }

  for (unsigned int n = 0; n < n_ratings; ++n)
    {
      unsigned int image = ratings[n].image;

      float decay = 
        exp (2.0 * 
             eta * 
             (sqrt (1.0 + lastimage[image]) - sqrt (1.0 + n_ratings)) /
             (n_ratings * priorlogbeta->sq_stddev));

      logbeta[image] = 
        decay * logbeta[image] + (1.0 - decay) * priorlogbeta->mean;
    }

  return qfunc (loga, n_workers, rank, gamma, logbeta, n_images,
                c, n_labels, d, pz, ratings, n_ratings,
                priorloga, priorgamma, &priorlogbeta->base, priorc, priord);
}
